# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/6 18:24 
@Author  : ZhangShenao 
@File    : 2.使用语义路由选择不同的Prompt模板.py 
@Desc    : 使用语义路由选择不同的Prompt模板

在RAG应用开发中,针对不同场景的问题,使用特定的prompt模板,效果一般都会比通用模板会好一些
例如在教培行业,针对不同的学科(语文、数学、英语),可以定义不同的Prompt模板
基于向量数据库的语义相似性检索能力,可以选择出与学生提问语义最相近的模板
这样RAG的回答效果会大大提升
"""
from typing import Dict, Any

import dotenv
from langchain.utils.math import cosine_similarity
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda

dotenv.load_dotenv()

# 为物理和数学两个学科,定义不同的Prompt模板
# 这两个模板的风格和回答思路都是不同的
physics_prompt = """你是一位非常聪明的物理教授。
你擅长以简洁易懂的方式回答物理问题。
当你不知道问题的答案时，你会坦率承认自己不知道。

这是一个问题：
{query}"""

math_prompt = """你是一位非常优秀的数学家。你擅长回答数学问题。
你之所以如此优秀，是因为你能将复杂的问题分解成多个小步骤。
并且回答这些小步骤，然后将它们整合在一起回来更广泛的问题。

这是一个问题：
{query}"""

# 创建Embedding模型,对Prompt进行嵌入
embeddings = ZhipuAIEmbeddings(model="embedding-3")
prompts = [physics_prompt, math_prompt]
prompt_embeddings = embeddings.embed_documents(prompts)


# 定义Prompt路由函数
def route_prompt(input: Dict[str, Any]) -> ChatPromptTemplate:
    """
    模板路由函数
    根据用户输入的提问,采用语义相似度检索,定位到语义上最相近的Prompt
    :param input: 用户的输入dict
    :return: 与用户输入在语义上最相近的Prompt模板
    """

    # 计算query的Embedding
    query_embedding = embeddings.embed_query(input["query"])

    # 使用LangChain内置的余弦相似度函数,计算query与每个Prompt之间的相似度
    similarity = cosine_similarity(X=[query_embedding], Y=prompt_embeddings)[0]
    selected = prompts[similarity.argmax()]
    print("使用数学Prompt" if selected == math_prompt else "使用物理模板")

    return ChatPromptTemplate.from_template(selected)


# 构建Chain
llm = ChatZhipuAI(model="glm-4-flash")
chain = (
        {"query": RunnablePassthrough()}
        | RunnableLambda(route_prompt)
        | llm
        | StrOutputParser()
)

print(chain.invoke("帮我简单介绍下黑洞是什么？"))
print("------------------------------------")
print(chain.invoke("能解释下导数吗？"))
